# -*- coding: utf-8 -*-
import numpy as np
from PyQt5.Qt import *
import sys
import json
import os

from resources import Ui_Form
from resources.utils import ImageShow, InferenceDetect, Segmentation, Roughness, SegMode, SegDetect


class MainPage(QWidget, Ui_Form):

    detect_signal = pyqtSignal(str)

    def __init__(self, parent=None, *args, **kwargs):
        super(MainPage, self).__init__(parent, *args, **kwargs)
        self.image_root = None
        self.label_image = None
        self.drag_flags = False
        self.start_locx, self.start_locy = None, None  # 设置拖拽事件
        self.ori_x, self.ori_y = None, None
        self.image_label_rect = None
        self.detect_result = None
        self.FINE_SEGMENT = 0    # 精分割模式

        self.setAttribute(Qt.WA_StyledBackground, True)
        self.setAttribute(Qt.WA_TranslucentBackground)
        self.setWindowFlags(Qt.FramelessWindowHint)

        self.timer = QTimer(self)
        self.timer.setSingleShot(True)
        self.timer.timeout.connect(lambda: self.notification_label.clear())
        self.setupUi(self)

        # 设置模型初始化
        self.detector = None
        self.segmentation_model = None
        self.seg_detect = None
        self.imageShow = ImageShow(label=self.image_label)

        # 信号
        self.min_btn.clicked.connect(self.showMinimized)
        self.max_restore_btn.clicked.connect(self.max_restore)
        self.shut_btn.clicked.connect(self.close)

        self.thre_spinBox.valueChanged[int].connect(lambda val: self.val_changed(val, 'thre_spinBox'))
        self.thre_slider.valueChanged[int].connect(lambda val: self.val_changed(val, 'thre_slider'))

    def max_restore(self):
        if self.max_restore_btn.isChecked():
            self.showMaximized()
            if self.image_root:
                pixmap = QPixmap(self.image_root).scaled(self.image_label.size(), aspectRatioMode=Qt.KeepAspectRatio)
                self.image_label.setPixmap(pixmap)
                self.image_label.setScaledContents(True)
                self.image_label.repaint()
        else:
            self.showNormal()
            if self.image_root:
                pixmap = QPixmap(self.image_root).scaled(QSize(self.imageShow.img_w, self.imageShow.img_h),
                                                         aspectRatioMode=Qt.KeepAspectRatio)
                self.image_label.setPixmap(pixmap)
                self.image_label.repaint()

    def set_detect_text(self, result, score_thr=0.7):
        bbox_result = result
        bboxes = np.vstack(bbox_result)
        scores = bboxes[:, -1]
        inds = scores > score_thr
        bboxes = bboxes[inds, :]
        labels = [
            np.full(bbox.shape[0], i, dtype=np.int32)
            for i, bbox in enumerate(bbox_result)
        ]
        labels = np.concatenate(labels)
        labels = labels[inds]
        self.textEdit.setText(f"已检测出{labels.shape[0]}处损伤:")
        for bbox, label in zip(bboxes, labels):
            if label == 0:
                labeltext = "scratch"
            else:
                labeltext = "pit"
            prob = format(bbox[-1]*100, ".2f")
            boxText = str(prob) + "%"
            self.textEdit.append("   "+labeltext + ": 可信度为 " + boxText)

    def open_file(self):
        self.image_label_rect = self.image_label.size()
        config_file = "resources/config/folder.json"   # 如需经常打开文件夹图片，请设置此路径
        config = json.load(open(config_file, "r", encoding="utf-8"))
        open_folder = config['open_folder']
        if not os.path.exists(open_folder):
            open_folder = os.getcwd()

        root, _ = QFileDialog.getOpenFileName(self, "打开图片", open_folder, "Image(*.png *.jpg *.JPEG)")
        # print(root)

        if not root:   # result：路径  格式raw
            return
        else:   # 使打开文件保留为上次打开文件的位置
            self.image_root = root
            self.notification_show(f"加载文件：{os.path.basename(root)}")
            config['open_folder'] = os.path.dirname(root)
            config_json = json.dumps(config, ensure_ascii=False, indent=2)
            with open(config_file, "w", encoding="utf-8") as f:
                f.write(config_json)
            # 显示文件
            self.imageShow.img = root
            self.imageShow.show_image()
            self.label_image = self.imageShow.image_data  # (480, 640, 3)

            self.segmentation_model = Segmentation(self.image_root)
            self.detector = InferenceDetect(self.image_root)
            self.seg_detect = SegDetect(self.image_root)

            self.comboBox.setCurrentIndex(0)
            self.textEdit.clear()

    def save_file(self):
        config_file = "resources/config/folder.json"   # 如需经常打开文件夹图片，请设置此路径
        config = json.load(open(config_file, "r", encoding="utf-8"))
        open_folder = config['open_folder']
        if not os.path.exists(open_folder):
            open_folder = os.getcwd()
        root, _ = QFileDialog.getSaveFileName(self, "保存图片", open_folder, "Image(*.png;*.jpg;*.JPEG)")

        if not root:
            return
        else:
            self.imageShow.save_img(root)  # 保存label中的图像
            self.notification_show(f"图像已保存：{os.path.basename(root)}")
            self.timer.start(5000)

    def notification_show(self, notify):
        self.notification_label.setText(notify)
        # 后面可设置显示时间 QTimer

    def detect_img(self):
        if self.image_root:
            # self.notification_show(f"正在检测：{os.path.basename(self.image_root)}")
            self.detector = InferenceDetect(self.image_root)
            detect_img, result = self.detector.image_detect()
            self.detect_result = result
            qImage = self.imageShow.ndarray2qimage(detect_img)
            self.image_label.setPixmap(QPixmap.fromImage(qImage))
            self.notification_show(f"检测完毕：{os.path.basename(self.image_root)}")
            self.set_detect_text(result, 0.7)

    def seg_img(self):
        if self.image_root:
            if not self.segmentation_model:
                self.segmentation_model = Segmentation(self.image_root)
            if not self.FINE_SEGMENT:
                # self.segmentation_model.img = self.image_root
                _, seg_img = self.segmentation_model.main_pipline()
                self.label_image = seg_img
            else:
                seg_img, seg = self.seg_detect.detect_seg(self.detect_result,
                                                          self.segmentation_model.seg_detect, self.detector.model)
                self.label_image = seg
            qImage = self.imageShow.ndarray2qimage(seg_img)
            self.image_label.setPixmap(QPixmap.fromImage(qImage))

    def b_channel(self):
        if self.image_root:
            if not self.segmentation_model:
                self.segmentation_model = Segmentation(self.image_root)
            self.segmentation_model.select_channel = self.segmentation_model.b

    def g_channel(self):
        if self.image_root:
            if not self.segmentation_model:
                self.segmentation_model = Segmentation(self.image_root)
            self.segmentation_model.select_channel = self.segmentation_model.g

    def r_channel(self):
        if self.image_root:
            if not self.segmentation_model:
                self.segmentation_model = Segmentation(self.image_root)
            self.segmentation_model.select_channel = self.segmentation_model.r

    def val_changed(self, val, name):
        if self.image_root:
            if not self.segmentation_model:
                self.segmentation_model = Segmentation(self.image_root)
            if name == 'thre_spinBox':
                self.thre_slider.setValue(val)
            elif name == 'thre_slider':
                self.thre_spinBox.setValue(val)
            self.segmentation_model.thre = val
            self.seg_img()

    def paint_mode(self, paint_area_mode):
        qImage = paint_area_mode.paint_area(self.image_label.mouse_rect.getCoords())
        self.label_image = paint_area_mode.processed_data[-1]
        self.image_label.setPixmap(QPixmap.fromImage(qImage))

    def select_mode(self, model):
        frame_img, self.label_image = model.select_area(self.image_label.mouse_rect.getCoords())
        self.image_label.setPixmap(QPixmap.fromImage(frame_img))

    def segment_mode_changed(self, index):
        if index == SegMode.SEGMENT:
            if self.image_root:
                self.imageShow.show_image()
                self.label_image = self.imageShow.image_data
            self.detect_btn.blockSignals(False)
            self.segmetation_btn.blockSignals(False)
            self.thre_slider.blockSignals(False)
            self.thre_spinBox.blockSignals(False)
            self.FINE_SEGMENT = 0
        elif index == SegMode.FINE_SEGMENT:
            self.detect_btn.blockSignals(True)
            self.FINE_SEGMENT = 1
            if self.image_root:
                self.imageShow.show_image()
                self.label_image = self.imageShow.image_data
                if not self.detect_result:
                    self.detect_img()
                if not self.segmentation_model:
                    self.segmentation_model = Segmentation(self.image_root)
                self.segmentation_model.select_channel = None
                if not self.detector:
                    self.detector = InferenceDetect(self.image_root)
                self.seg_detect = SegDetect(self.image_root)
        elif index == SegMode.PAINT_AREA:
            self.detect_btn.blockSignals(True)   # 暂时阻断其他功能
            self.segmetation_btn.blockSignals(True)
            self.thre_slider.blockSignals(True)
            self.thre_spinBox.blockSignals(True)
            if self.image_root:
                self.imageShow.show_image()
                self.label_image = self.imageShow.image_data
                paint_area_mode = SegMode(self.image_root, self.image_label.geometry(),
                                          self.image_label.pixmap().rect(), self.imageShow.scale_ratio)
                self.image_label.mouse_move_signal.connect(lambda model=paint_area_mode: self.paint_mode(model))
        elif index == SegMode.SELECT_AREA:
            self.detect_btn.blockSignals(True)
            self.segmetation_btn.blockSignals(True)
            self.thre_slider.blockSignals(True)
            self.thre_spinBox.blockSignals(True)
            if self.image_root:
                self.imageShow.show_image()
                self.label_image = self.imageShow.image_data
                paint_area_mode = SegMode(self.image_root, self.image_label.geometry(),
                                          self.image_label.pixmap().rect(), self.imageShow.scale_ratio)
                self.image_label.mouse_move_signal.connect(lambda model=paint_area_mode: self.select_mode(model))

    def calculate_roughness(self):
        if self.image_root:
            # image = self.imageShow.qimage2ndarray()  # 直接获取label内图像计算，不是很准
            image = self.imageShow.rgb2gray(self.label_image)
            calculator = Roughness(image)
            self.roughness_le.setText(str(calculator.Sa))

    def mousePressEvent(self, evt) -> None:
        if evt.button() == Qt.LeftButton:
            self.drag_flags = True
            self.start_locx, self.start_locy = evt.globalX(), evt.globalY()
            self.ori_x, self.ori_y = self.x(), self.y()

    def mouseMoveEvent(self, evt) -> None:
        if self.drag_flags:
            self.move(self.ori_x + evt.globalX() - self.start_locx, self.ori_y + evt.globalY() - self.start_locy)

    def mouseReleaseEvent(self, evt) -> None:
        self.drag_flags = False

    @staticmethod
    def signal_reset(used_widget):
        used_widget.blockSignals(True)
        used_widget.clear()
        used_widget.blockSignals(False)


if __name__ == '__main__':

    app = QApplication(sys.argv)
    window = MainPage()

    window.show()
    sys.exit(app.exec_())